"""
使用typing包的Optional，明确标识可能为 None 的值：提高代码可读性，提醒处理 None 情况。
"""

from io import BytesIO
from uuid import uuid4

from langgraph.checkpoint.memory import MemorySaver
from langgraph.graph.message import add_messages
from langchain_core.messages import HumanMessage, AIMessage
from typing import TypedDict, Optional, List
from langgraph.graph.state import CompiledStateGraph

from app.api.human_input_handle import get_thread_identifier
from app.functions.generate_report import generate_report_node
from app.utils.llm_client import LLMClient
from app.utils.mermaid_code_util import visualize_with_mermaid_code
from app.utils.knowledge_base import KnowledgeBase
from app.models import SessionState

from app.functions.classify_issue_or_require import classify_user_input, InputType, classify_user_input_node
from app.functions.require_clarification import require_clear_or_not, require_clarify, ClearOrNot, \
    require_clear_or_not_node
from app.functions.require_identificaton_and_analysis import FieldAvailableOrNot, field_availability_check
from app.functions.require_reconfirm_extract import require_reconfirm, ChangeOrNot
from app.functions.require_identificaton_and_analysis import field_availability_check_node
#20250820新增
from langgraph.graph import StateGraph, START, END
from app.functions.require_reconfirm_extract import require_reconfirm_extract_node


class DataRequirementAnalysisWorkflowLogic:
    """
    该类为总体逻辑工作流，指引后续的工作流开发工作，实际使用工作流以DataRequirementAnalysisWorkflow为准

    """
    def __init__(
        self,
            llm_client: Optional[LLMClient] = None,
            knowledge_base: Optional[KnowledgeBase] = None
    ):
        """初始化工作流,选择大模型需要调用的大模型、知识库

        Args:
            llm_client: 大模型客户端
            knowledge_base: 知识库客户端
        """
        self.llm_client = llm_client or LLMClient()
        self.knowledge_base = knowledge_base or KnowledgeBase()

        self.graph = self._build_workflow_graph()


    def _build_workflow_graph(self) -> CompiledStateGraph:
        """
        构建需求分析智能体的总体graph，并编译
        """
        workflow = StateGraph(SessionState)

        workflow.add_node("start", self._start_node)
        # 判断为数据问题或数据需求的节点，后接条件边xxxx
        workflow.add_node("classify_user_input", self._classify_user_input)
        # 处理数据问题的节点，目前为仅打印信息的空节点
        workflow.add_node("issue_process", self._issue_process)
        # 判断需求是否清晰的节点，后接条件边xxxx
        workflow.add_node("require_clear_or_not", self._require_clear_or_not)
        # 需求澄清节点，后接普通边xxxx
        workflow.add_node("require_clarify", self._require_clarify)
        # 判别字段是否具备，并维护field_json的节点，后接条件边xxx
        workflow.add_node("require_transfer", self._require_transfer)
        # 产出涉加工方案的节点，后接普通边xxx
        workflow.add_node("deliver_process_solution", self._deliver_process_solution)
        # 要求用户二次确认不具备字段的节点，如果用户不删除或修改不具备的字段，走采集向涉采集方案的节点，后接条件边
        workflow.add_node("change_require_reconfirm", self._change_require_reconfirm)
        # 产出涉采集方案的节点，后接普通边
        workflow.add_node("deliver_collect_solution", self._deliver_collect_solution)
        # 产出最终报告的节点，后接结束边
        workflow.add_node("deliver_final_solution_report", self._deliver_final_solution_report)
        workflow.add_node("end", self._end_node)
        #备用node，处理错误信息
        workflow.add_node("error", self._error_node)

        # 为图补全所有边
        self._add_workflow_edges(workflow)
        # 设置入口点
        workflow.set_entry_point("start")

        return workflow.compile()

    def _add_workflow_edges(self, workflow: StateGraph):
        """添加工作流边

        Args:
            workflow: StateGraph实例
        """
        # 从start到issue_or_require
        workflow.add_edge("start", "classify_user_input")
        # 条件边，分别指向问题处理节点和需求澄清节点
        workflow.add_conditional_edges("classify_user_input",
                                       self._issue_require_router,
                                       {
                                           "issue": "issue_process",
                                           "requirement": "require_clear_or_not",
                                           "error": "error"
                                       }
                                       )
        workflow.add_conditional_edges("require_clear_or_not",
                                       self._require_clear_or_not_router,
                                       {
                                           "clear": "require_transfer",
                                           "not_clear": "require_clarify",
                                           "error": "error"
                                       })
        workflow.add_conditional_edges("require_transfer",
                                       self._field_availability_router,
                                       {
                                           "available": "deliver_process_solution",
                                           "not_available": "change_require_reconfirm",
                                           "error": "error"
                                       })
        workflow.add_edge("deliver_process_solution", "deliver_final_solution_report")
        workflow.add_conditional_edges("change_require_reconfirm",
                                       self._change_unavaliable_field_or_not_router,
                                       {
                                           "not_change": "deliver_collect_solution",
                                           "change": "require_transfer",
                                           "error": "error"
                                       })
        workflow.add_edge("deliver_process_solution", "deliver_final_solution_report")
        workflow.add_edge("deliver_collect_solution", "deliver_final_solution_report")
        workflow.add_edge("deliver_final_solution_report","end")
        workflow.add_edge("error", END)
        workflow.add_edge("end", END)

    """
    以下是nodes的具体定义
    """
    def _start_node(self, state: SessionState) -> SessionState:

        try:
            # TODO：提示当前节点状态
            return state

        except Exception:
            # TODO: 为日志添加错误信息
            # TODO: 为state添加错误信息
            return state

    def _classify_user_input(self, state: SessionState) -> SessionState:

        try:
            # TODO：提示当前节点状态

            # 获取分类结果
            classification_result = classify_user_input(
                state.original_input, self.llm_client
            )
            # TODO：用classification_result更新state特定值的状态

            # TODO：记录并更新日志
            return state

        except Exception:
            # TODO: 为日志添加错误信息
            # TODO: 为state添加错误信息
            return state

    def _issue_process(self, state: SessionState) -> SessionState:
        try:
            # TODO：提示当前节点状态
            # TODO: 获取分类结果,v1.1演进
            # TODO：用classification_result更新state状态
            # TODO：记录并更新日志
            return state

        except Exception:
            # TODO: 为日志添加错误信息
            # TODO: 为state添加错误信息
            return state

    def _require_clear_or_not(self, state: SessionState) -> SessionState:

        try:
            # TODO：提示当前节点状态

            # 获取分类结果
            clear_or_not = require_clear_or_not(
                state.require_json, self.llm_client
            )
            # TODO：用clear_or_not更新state状态

            # TODO：记录并更新日志
            return state

        except Exception:
            return state

    def _require_clarify(self, state: SessionState) -> SessionState:
        try:
            # TODO：提示当前节点状态

            # 获取分类结果
            clear_or_not = require_clarify(
                state.clarify_input
            )
            # TODO：用clear_or_not更新state状态

            # TODO：记录并更新日志
            return state

        except Exception:
            return state

    def _require_transfer(self, state: SessionState) -> SessionState:
        try:
            # TODO：提示当前节点状态

            # 获取分类结果,field_availability包含两个信息，一个是而分类的是否完全具备，一个是详细的字段级具备情况说明
            field_availability = field_availability_check(
                state.require_json, self.llm_client
            )
            # TODO：用field_availability更新state状态

            # TODO：记录并更新日志
            return state

        except Exception:
            return state

    def _deliver_process_solution(self, state: SessionState) -> SessionState:
        try:
            # TODO：提示当前节点状态

            # 获取分类结果,field_availability包含两个信息，一个是而分类的是否完全具备，一个是详细的字段级具备情况说明
            # process_solution = deliver_process_solution(
            #     state.field_json, self.llm_client
            # )
            # TODO：用process_solution更新state状态

            # TODO：记录并更新日志
            return state

        except Exception:
            return state

    def _change_require_reconfirm(self, state: SessionState) -> SessionState:
        try:
            # TODO：提示当前节点状态

            # 获取分类结果,field_availability包含两个信息，一个是而分类的是否完全具备，一个是详细的字段级具备情况说明
            reconfirmed_require = require_reconfirm(
                state.field_json, state.require_json, self.llm_client
            )
            # TODO：用reconfirmed_require更新state状态

            # TODO：记录并更新日志
            return state

        except Exception:
            return state

    def _deliver_collect_solution(self, state: SessionState) -> SessionState:
        try:
            # TODO：提示当前节点状态

            # 获取分类结果,field_availability包含两个信息，一个是而分类的是否完全具备，一个是详细的字段级具备情况说明
            # collect_solution = deliver_collect_solution(
            #     state.field_json, self.llm_client
            # )
            # TODO：用process_solution更新state状态

            # TODO：记录并更新日志
            return state

        except Exception:
            return state

    def _deliver_final_solution_report(self, state: SessionState) -> SessionState:
        try:
            # TODO：提示当前节点状态

            # final_solution_report = deliver_final_solution(
            #     state.process_solution, state.require_json, self.llm_client
            # )
            # TODO：用process_solution更新state状态

            # TODO：记录并更新日志
            return state

        except Exception:
            return state

    def _end_node(self, state: SessionState) -> SessionState:
        try:
            # TODO：提示当前节点状态

            # TODO：直接更新state状态，可能需要更新工作流状态

            # TODO：记录并更新日志
            return state
        except Exception:
            return state

    def _error_node(self, state: SessionState) -> SessionState:
        """错误处理节点"""
        # TODO: 进入错误状态提示，并更新日志

        # TODO：更新工作流状态为error？

        # 生成错误报告
        try:
            # TODO：更新state中的error处理，并生成错误报告
            return state

        except Exception:
            # TODO：处理如生成错误报告失败的错误
            return state
    """
    以下是routers
    """
    def _issue_require_router(self, state: SessionState) -> str:
        if state.input_type == InputType.REQUIREMENT:
            return "requirement"
        if state.input_type == InputType.ISSUE:
            return "issue"
        else:
            return "error"

    def _require_clear_or_not_router(self, state: SessionState) -> str:
        if state.require_clearness == ClearOrNot.CLEAR:
            return "clear"
        if state.input_type == ClearOrNot.NOT_CLEAR:
            return "not_clear"
        else:
            return "error"

    def _field_availability_router(self, state: SessionState) -> str:
        if state.field_availability == FieldAvailableOrNot.AVAILABLE:
            return "available"
        if state.field_availability == FieldAvailableOrNot.NOT_AVAILABLE:
            return "not_available"
        else:
            return "error"

    def _change_unavaliable_field_or_not_router(self, state: SessionState) -> str:
        if state.change_or_not == ChangeOrNot.CHANGE:
            return "change"
        if state.change_or_not == ChangeOrNot.NOT_CHANGE:
            return "not_change"
        else:
            return "error"



class DataRequirementAnalysisWorkflow:
   
    def __init__(self):
        self.graph = None
        self.memory = None

    def build(self):
        builder = StateGraph(SessionState)
        builder.add_node("human_initial", lambda state: state)  # 新增占位节点,以免用户直接进入第一个节点
        builder.add_node("classify_user_input", lambda state,config:classify_user_input_node(state, config))
        builder.add_node("require_clear_or_not", lambda state,config:require_clear_or_not_node(state, config))
        builder.add_node("human_clarify_requirement",lambda state: state)
        builder.add_node("field_availability_check", lambda state,config:field_availability_check_node(state, config))
        builder.add_node("require_reconfirm_extract", lambda state,config:require_reconfirm_extract_node(state, config))
        builder.add_node("human_require_confirm", lambda state:state)
        builder.add_node("generate_report",  lambda state,config:generate_report_node(state, config)) # 生成报告的功能节点

        builder.add_edge(START, "human_initial")
        builder.add_edge("human_initial", "classify_user_input")

        # 分类用户输入后的条件跳转
        builder.add_conditional_edges(
            "classify_user_input",
            lambda x: x["next"],
            {"require_clear_or_not": "require_clear_or_not", END:END}
        )

        # 需求明确性检查后的跳转
        builder.add_conditional_edges(
            "require_clear_or_not",
            lambda x: x["next"],
            {"human_clarify_requirement": "human_clarify_requirement",
             "field_availability_check": "field_availability_check"}
        )
        builder.add_edge("human_clarify_requirement", "require_clear_or_not")


        builder.add_conditional_edges(
            "field_availability_check",
            lambda x: x["next"],  # 这里的x["next"]可以是"human_require_confirm"或"generate_report"
            {"human_require_confirm": "human_require_confirm",
             "generate_report": "generate_report"}
        )

        # 形成完整循环路径：human_require_confirm -> require_reconfirm_extract -> field_availability_check
        builder.add_edge("human_require_confirm", "require_reconfirm_extract")
        builder.add_edge("require_reconfirm_extract", "field_availability_check")
        builder.add_edge("generate_report",END)

        self.memory = MemorySaver()
        self.graph = builder.compile(
            checkpointer=self.memory,
            interrupt_before=["classify_user_input", "human_clarify_requirement", "human_require_confirm"]   # 断点位置，
        )
        print(visualize_with_mermaid_code(self.graph))
        return self.graph


    def run(self, state: Optional[SessionState] = None, config:Optional[dict] = None):
        if self.graph is None:
            self.graph=self.build()

        # 初始化状态  
        if state is None:
            state: SessionState={
                "original_input": None,
                "input_type": None,
                "require_clearness": None,
                "clarify_input": None,
                "require_json": None,
                "field_availability": None,
                "field_json": None,
                "process_solution": None,
                "collect_solution": None,
                "change_or_not": None,
                "next": None,  # 初始不指定下一步，由流程逻辑决定
                "messages": [] ,# 消息列表初始化为空列表
                "user_initial_input": None ,  #
                "missing_fields": None,  # 缺失字段信息
                "user_require":[],
                "user_loop_count":0, #用户多次更新字段需求最大次数，暂定最大3次
                "user_satisfaction_intention":0, #用户满意程度意图标识，0为不满意，1为满意
                "user_update_fileds":None ,#用于存储提取出来的字段信息


            }
        thread_id = get_thread_identifier()
        if config is None:
            config = {"configurable": {"thread_id": thread_id}}
        # 初始化状态
        self.graph.update_state(config, state)
        # 运行流程直到断点
        current_state = state
        while True:
            # 运行流程直到断点
            for event in self.graph.stream(None, config, stream_mode="values"):
                current_state = event

            # 显示助手消息
            if current_state.get("messages"):
                print(f"助手: {current_state['messages'][-1].content}")

            # 检查是否结束
            if current_state.get("next") == END:
                print("会话结束。")
                break

            # 断点时的输入
            user_input = input("请输入您的内容: ")
            current_state["messages"]=add_messages(current_state["messages"], [HumanMessage(content=user_input)])
            self.graph.update_state(config, current_state)


class DataRequirementAnalysisWorkflow2:
    """
    开发过程中实际使用的工作流，不断补全工作流的节点
    """
    builder = StateGraph(SessionState)
    builder.add_node("human_initial", lambda state: state)  # 新增占位节点,以免用户直接进入第一个节点
    builder.add_node("classify_user_input", classify_user_input_node)
    builder.add_node("require_clear_or_not", require_clear_or_not_node)
    builder.add_node("human_clarify_requirement",lambda state: state)
    builder.add_node("field_availability_check", field_availability_check_node)
    builder.add_node("require_reconfirm_extract", require_reconfirm_extract_node)
    builder.add_node("human_require_confirm", lambda state:state)
    builder.add_node("generate_report", generate_report_node) # 生成报告的功能节点

    builder.add_edge(START, "human_initial")
    builder.add_edge("human_initial", "classify_user_input")

    # 分类用户输入后的条件跳转
    builder.add_conditional_edges(
        "classify_user_input",
        lambda x: x["next"],
        {"require_clear_or_not": "require_clear_or_not", END:END}
    )

    # 需求明确性检查后的跳转
    builder.add_conditional_edges(
        "require_clear_or_not",
        lambda x: x["next"],
        {"human_clarify_requirement": "human_clarify_requirement",
         "field_availability_check": "field_availability_check"}
    )
    builder.add_edge("human_clarify_requirement", "require_clear_or_not")


    builder.add_conditional_edges(
        "field_availability_check",
        lambda x: x["next"],  # 这里的x["next"]可以是"human_require_confirm"或"generate_report"
        {"human_require_confirm": "human_require_confirm",
         "generate_report": "generate_report"}
    )

    # 形成完整循环路径：human_require_confirm -> require_reconfirm_extract -> field_availability_check
    builder.add_edge("human_require_confirm", "require_reconfirm_extract")
    builder.add_edge("require_reconfirm_extract", "field_availability_check")
    builder.add_edge("generate_report",END)

    memory = MemorySaver()
    graph = builder.compile(
        checkpointer=memory,
        interrupt_before=["classify_user_input", "human_clarify_requirement", "human_require_confirm"]  # 断点位置
    )
    print(visualize_with_mermaid_code(graph))
   # visualize_graph(graph)







